#include "kernel_operator.h"
#include "arg_max_perf.h"
// #include "kernel_float.h"
// #include "impl/"
#include <cstdint>
using namespace AscendC;
// #define K_MAX_SHAPE_DIM 0


// constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue
template<typename T>class KernelArgMaxValueSpec{
public:
    __aicore__ inline KernelArgMaxValueSpec(){}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR indice, GM_ADDR values,
     // uint32_t totalLength, uint32_t resLength, uint32_t stride,
     // uint32_t iterStep, uint32_t nIter, int64_t axes, uint32_t dims, 
     uint32_t axesDim,
     uint32_t smallDataNum,uint32_t bigDataNum,uint32_t tailBlockNum, TPipe* pipeIn){
        // this->totalLength = totalLength;
        // this->resLength = resLength;
        // this->stride = stride;
        // this->iterStep = iterStep;
        // this->nIter = nIter;
        // this->axes = axes;
        // this->dims = dims;
        this->axesDim = axesDim;
        this->roundDim = (axesDim + 63)/64*64;
        pipe = pipeIn; 
        // this->roundDim = 
        uint32_t coreNum = AscendC::GetBlockIdx();
        uint32_t globalIndex = bigDataNum * coreNum;
        if (coreNum < tailBlockNum) {
            this->coreDataNum = bigDataNum;
        } else {
            this->coreDataNum = smallDataNum;
            globalIndex -= (bigDataNum - smallDataNum)*(coreNum - tailBlockNum);
        }
        this->roundcoreDataNum = (this->coreDataNum + 7)/8*8;
        this->globalIndex = globalIndex;
        // printf("global: %d coreDataNum:%d\n", globalIndex, coreDataNum);
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + globalIndex*axesDim, this->coreDataNum*axesDim *sizeof(DTYPE_X));
        indiceGm.SetGlobalBuffer((__gm__ DTYPE_INDICE*)indice + globalIndex, this->coreDataNum * sizeof(DTYPE_INDICE));
        valuesGm.SetGlobalBuffer((__gm__ DTYPE_VALUES*)values + globalIndex, this->coreDataNum * sizeof(DTYPE_VALUES));
        
        // this->loopCount = this->coreDataNum;
        
        pipe->InitBuffer(inQueueX, this->roundDim*this->coreDataNum*sizeof(DTYPE_X));
        // pipe->InitBuffer(outQueue, 32*this->coreDataNum*sizeof(DTYPE_X));

        pipe->InitBuffer(outQueueValues, this->roundcoreDataNum*sizeof(DTYPE_X));
        pipe->InitBuffer(outQueueIndice, this->roundcoreDataNum*sizeof(int32_t));
        pipe->InitBuffer(indexQueue, this->roundcoreDataNum*sizeof(DTYPE_X));
        pipe->InitBuffer(WorkQueue, this->roundcoreDataNum*sizeof(DTYPE_X));
        pipe->InitBuffer(WorkQueue2, this->roundcoreDataNum*sizeof(DTYPE_X));

    }
    __aicore__ inline void process() {
        copyIn();
        compute();
        // copyOut();
    }
    __aicore__ inline void process1() {
        copyIn1();
        compute1();
        copyOut();
    }

    __aicore__ inline void copyIn1() {
        auto xLocal = inQueueX.Get<DTYPE_X>();
        AscendC::DataCopyExtParams copyParams{static_cast<uint16_t>(coreDataNum),  static_cast<uint32_t>(axesDim*4), 0, 0,0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        // AscendC::DataCopyPadExtParams<DTYPE_X> padParams{false, 0, static_cast<uint8_t>(roundDim - axesDim), 0};
        AscendC::DataCopyPadExtParams<DTYPE_X> padParams{true, 0, 0, DTYPE_X(-1000)};

        // AscendC::DataCopyExtParams copyParams{static_cast<uint16_t>(coreDataNum),  static_cast<uint32_t>(axesDim*4), 0,  (64-(axesDim+7)/8),0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        // AscendC::DataCopyPadExtParams<DTYPE_X> padParams{true, 0, 0, DTYPE_X(-1000)};
        AscendC::DataCopyPad(xLocal, xGm, copyParams,padParams);
        xLocal.GetValue(0);
        // inQueueX.EnQue(xLocal);

    }
    __aicore__ inline void copyIn() {
        auto xLocal = inQueueX.Get<DTYPE_X>();
        // AscendC::DataCopyExtParams copyParams{static_cast<uint16_t>(coreDataNum),  static_cast<uint32_t>(axesDim*4), 0, 0,0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        // AscendC::DataCopyPadExtParams<DTYPE_X> padParams{false, 0, static_cast<uint8_t>(roundDim - axesDim), 0};
        AscendC::DataCopyExtParams copyParams{static_cast<uint16_t>(coreDataNum),  static_cast<uint32_t>(axesDim*4), 0,  (64-(axesDim+7)/8),0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPadExtParams<DTYPE_X> padParams{true, 0,  static_cast<uint8_t>(8 - axesDim%8), DTYPE_X(-1000)};
        AscendC::DataCopyPad(xLocal, xGm, copyParams,padParams);
        xLocal.GetValue(0);
        // inQueueX.EnQue(xLocal);
        // for (int i = 0; i < 512; i++) xLocal.GetValue(i);
    }
__aicore__ inline void compute1() {
        AscendC::LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        // AscendC::LocalTensor<DTYPE_X> dstLocal = outQueue.Get<DTYPE_X>();
        // AscendC::LocalTensor<DTYPE_X> workLocal = WorkQueue.Get<DTYPE_X>();
        // AscendC::LocalTensor<DTYPE_X> workLocal2 = WorkQueue2.Get<DTYPE_X>();
        auto valueLocal = outQueueValues.Get<DTYPE_X>();
        auto indexLocal = outQueueIndice.Get<int32_t>();
        AscendC::WholeReduceMax<DTYPE_X>(xLocal, xLocal, 64, coreDataNum*(roundDim/64), 1, 1, 8, AscendC::ReduceOrder::ORDER_VALUE_INDEX);
        // AscendC::WholeReduceMax<DTYPE_X>(workLocal2, dstLocal, 16, coreDataNum, 1, 1, 2, AscendC::ReduceOrder::ORDER_VALUE_INDEX);
        for (int i = 0; i < coreDataNum; i++) {
            DTYPE_X maxValue = xLocal.GetValue(i*2);
            DTYPE_X maxIndex = xLocal.GetValue(i*2+1);
            // uint32_t blockIndex = *reinterpret_cast<uint32_t *>(&maxIndex);
            // DTYPE_X innerIndex = dstLocal.GetValue(i*16  + blockIndex+1);
            // // printf("block %d %d\n", blockIndex, innerIndex);
            uint32_t realIndex = *reinterpret_cast<uint32_t *>(&maxIndex);
            valueLocal.SetValue(i, maxValue);
            indexLocal.SetValue(i, realIndex);
            // printf("%f %d\n", maxValue, blockIndex*32 + realInnerIndex);
        }
    }
    __aicore__ inline void compute() {
        AscendC::LocalTensor<DTYPE_X> xLocal = inQueueX.Get<DTYPE_X>();
        AscendC::LocalTensor<int32_t> index = indexQueue.Get<int32_t>();
        AscendC::LocalTensor<DTYPE_X> workLocal = WorkQueue.Get<DTYPE_X>();
        AscendC::LocalTensor<DTYPE_X> workLocal2 = WorkQueue2.Get<DTYPE_X>();
        auto valueLocal = outQueueValues.Get<DTYPE_X>();
        auto indexLocal = outQueueIndice.Get<int32_t>();
        AscendC::WholeReduceMax<DTYPE_X>(xLocal, xLocal, 64, coreDataNum*(roundDim/64), 1, 1, 8, AscendC::ReduceOrder::ORDER_VALUE_INDEX);
        // AscendC::PipeBarrier<PIPE_V>();
        // AscendC::WholeReduceMax<DTYPE_X>(workLocal2, xLocal, 16, coreDataNum, 1, 1, 2, AscendC::ReduceOrder::ORDER_VALUE_INDEX);
        AscendC::WholeReduceMax<DTYPE_X>(workLocal2, xLocal, 16, coreDataNum, 1, 1, 2, AscendC::ReduceOrder::ORDER_ONLY_INDEX);

        AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(coreDataNum * 4), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
//         valueLocal.GetValue(0);

        auto workLocal2Int = workLocal2.ReinterpretCast<int32_t>();
        auto workLocalInt = workLocal.ReinterpretCast<int32_t>();
//         // AscendC::PipeBarrier<PIPE_V>();
        AscendC::CreateVecIndex(workLocalInt, (int32_t)0, coreDataNum);
        AscendC::Muls(workLocalInt, workLocalInt, int32_t(16) , coreDataNum);

        AscendC::Add(workLocalInt,workLocalInt, workLocal2Int, coreDataNum);
        AscendC::Muls(workLocalInt, workLocalInt, int32_t(4) , coreDataNum);
        Gather(valueLocal, xLocal, workLocalInt.ReinterpretCast<uint32_t>(), (uint32_t)0,coreDataNum);
        AscendC::Adds(workLocalInt, workLocalInt, int32_t(4) , coreDataNum);
        Gather(indexLocal, xLocal.ReinterpretCast<int32_t>(), workLocalInt.ReinterpretCast<uint32_t>(), (uint32_t)0,coreDataNum);

        AscendC::Muls(workLocal2Int, workLocal2Int, int32_t(32) , coreDataNum);
        AscendC::Add(indexLocal,indexLocal, workLocal2Int, coreDataNum);
        indexLocal.GetValue(0);

        AscendC::DataCopyPad(indiceGm, indexLocal, copyParams); // 从GM->VECIN搬运40Bytes
        AscendC::DataCopyPad(valuesGm, valueLocal, copyParams); // 从GM->VECIN搬运40Bytes

    }
    __aicore__ inline void copyOut() {
        // LocalTensor<DTYPE_X>dstLocal = outQueue.DeQue<DTYPE_X>();
        auto valueLocal = outQueueValues.Get<DTYPE_X>();
        auto indexLocal = outQueueIndice.Get<int32_t>();
        // for (int i = 0; i < coreDataNum; i++) {
        //     printf("%f %d\n", valueLocal.GetValue(i), indexLocal.GetValue(i));
        // }
        // valueLocal.GetValue(coreDataNum-1);
        // indexLocal.GetValue(coreDataNum-1);

        // printf("%f %d\n", valueLocal.GetValue(0), indexLocal.GetValue(0));

        // auto dstLocalRein = dstLocal.ReinterpretCast<float>();
        // for (int i = 0; i < coreDataNum; i++) {
        //     float maxValue = dstLocalRein.GetValue(i*roundDim/64*2);
        //     float maxIndex = dstLocalRein.GetValue(i*roundDim/64*2+1);
        //     uint32_t realIndex = *reinterpret_cast<uint32_t *>(&maxIndex);
        //     for (int j = 1; j < (roundDim/64); j++) {
        //         float now =  dstLocalRein.GetValue((i*roundDim/64 + j)*2);
        //         float nowIndex = dstLocalRein.GetValue((i*roundDim/64 + j)*2 + 1);
        //         uint32_t realNowIndex = *reinterpret_cast<uint32_t *>(&nowIndex) + 64*j;
        //         if (maxValue > now) {
        //             maxValue = now;
        //             realIndex = realNowIndex;
        //         }
        //     }
        //     indexLocal.SetValue(i, realIndex);
        //     valueLocal.SetValue(i, maxValue);
        // }
        // DataCopy(valuesGm, valueLocal, roundcoreDataNum);
        // DataCopy(indiceGm, indexLocal, roundcoreDataNum);
        // outQueue.FreeTensor(dstLocal);
        AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(coreDataNum * 4), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位

        AscendC::DataCopyPad(valuesGm, valueLocal, copyParams); // 从GM->VECIN搬运40Bytes
        AscendC::DataCopyPad(indiceGm, indexLocal, copyParams); // 从GM->VECIN搬运40Bytes
    }

private:
    int32_t coreDataNum;
    // uint32_t totalLength;
    uint32_t globalIndex;
    // uint32_t resLength;
    // uint32_t stride;
    // uint32_t iterStep;
    // uint32_t nIter;
    int64_t axes;
    uint32_t dims;
    uint32_t axesDim;
    uint32_t roundDim,roundcoreDataNum;
    // uint32_t loopCount;
    AscendC::TPipe *pipe;
    AscendC::GlobalTensor<DTYPE_X> xGm;
    AscendC::GlobalTensor<DTYPE_INDICE> indiceGm;
    AscendC::GlobalTensor<DTYPE_VALUES> valuesGm;
    // TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    // TQue<QuePosition::VECOUT, BUFFER_NUM>  ;
    // TBuf<QuePosition::VECCALC>outQueueValues, outQueueIndice,inQueueX;
    TBuf<QuePosition::VECCALC>WorkQueue,outQueue,WorkQueue2, outQueueValues, outQueueIndice,inQueueX,indexQueue;
};

template<typename T>class KernelArgMaxValue{
public:
    __aicore__ inline KernelArgMaxValue(){}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR indice, GM_ADDR values,
     uint32_t totalLength, uint32_t resLength, uint32_t stride,
     uint32_t iterStep, uint32_t nIter, int64_t axes, uint32_t dims, uint32_t axesDim,
     uint32_t smallDataNum,uint32_t bigDataNum,uint32_t tailBlockNum, TPipe* pipeIn){
        this->totalLength = totalLength;
        this->resLength = resLength;
        this->stride = stride;
        this->iterStep = iterStep;
        this->nIter = nIter;
        this->axes = axes;
        this->dims = dims;
        this->axesDim = axesDim;
        pipe = pipeIn; 

        uint32_t coreNum = AscendC::GetBlockIdx();
        uint32_t globalIndex = bigDataNum * coreNum;
        if (coreNum < tailBlockNum) {
            this->coreDataNum = bigDataNum;
        } else {
            this->coreDataNum = smallDataNum;
            globalIndex -= (bigDataNum - smallDataNum)*(coreNum - tailBlockNum);
        }
        this->globalIndex = globalIndex;
        printf("global: %d coreDataNum:%d\n", globalIndex, coreDataNum);
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x, this->totalLength *sizeof(DTYPE_X));
        indiceGm.SetGlobalBuffer((__gm__ DTYPE_INDICE*)indice + globalIndex, this->coreDataNum * sizeof(DTYPE_INDICE));
        valuesGm.SetGlobalBuffer((__gm__ DTYPE_VALUES*)values + globalIndex, this->coreDataNum * sizeof(DTYPE_VALUES));
        
        this->loopCount = this->coreDataNum;
        pipe->InitBuffer(Queue_tempIndex,  2 * sizeof(int32_t));
        pipe->InitBuffer(Queue_tempValues, 2 * sizeof(DTYPE_X));

        pipe->InitBuffer(inQueueX, BUFFER_NUM, this->axesDim*8*sizeof(DTYPE_X));
        // // pipe->InitBuffer(inQueueFullX, BUFFER_NUM, this->axesDim*8*sizeof(DTYPE_X));

        // pipe->InitBuffer(outQueueValues, BUFFER_NUM, this->axesDim*sizeof(DTYPE_X));
        // pipe->InitBuffer(outQueueIndice, BUFFER_NUM, this->axesDim*sizeof(int32_t));
        pipe->InitBuffer(computeBuf, 2*sizeof(DTYPE_X));
        pipe->InitBuffer(WorkQueue, this->axesDim*8*sizeof(DTYPE_X));


        this->tempIndex = Queue_tempIndex.Get<int32_t>();
        // this->xbuf = xFullBuf.Get<DTYPE_X>();
        // Duplicate(this->tempIndex, int32_t(0), this->stride);

        this->tempValue = Queue_tempValues.Get<DTYPE_X>();
        // Duplicate(this->tempValue, DTYPE_X(0),this->stride);
    }
    __aicore__ inline void process() {
        for (int i = globalIndex; i < globalIndex + loopCount; i++) {
            printf("i: %d\n",i);
            copyIn(i/stride, i%stride);
            compute(i);
            copyOut(i-globalIndex);
        }
    }

    __aicore__ inline void copyIn(int progress, int now) {
        LocalTensor<DTYPE_X>xLocal = inQueueX.AllocTensor<DTYPE_X>();
        AscendC::DataCopyExtParams copyParams{static_cast<uint16_t>(this->axesDim),  4, (this->stride-1) * 4, 0,0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPadExtParams<DTYPE_X> padParams{false, 0, 0, 0};
        // int x = process - process/stride;
        AscendC::DataCopyPad(xLocal, xGm[progress*iterStep + now], copyParams, padParams); // 从GM->VECIN搬运40Bytes
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void compute(int progress) {
        AscendC::LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        AscendC::LocalTensor<DTYPE_X> dstLocal = computeBuf.Get<DTYPE_X>();
        AscendC::LocalTensor<DTYPE_X> workLocal = WorkQueue.Get<DTYPE_X>();
        AscendC::ReduceMax(dstLocal ,xLocal, workLocal, this->axesDim*8, true);
        DTYPE_X maxValue = dstLocal.GetValue(0);
        DTYPE_X maxIndex = dstLocal.GetValue(1);
        uint32_t realIndex = *reinterpret_cast<uint32_t *>(&maxIndex);
        realIndex /= 8;
        tempIndex.SetValue(0, realIndex);
        tempValue.SetValue(0, maxValue);
        inQueueX.FreeTensor(xLocal);
    }

    __aicore__ void copyOut(int32_t progress) {
        AscendC::DataCopyExtParams copyParams{1, 1 * 4, 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        // AscendC::DataCopyPadExtParams<DTYPE_X> padParams{false, 0, 0, 0};
        // AscendC::DataCopyPadExtParams<int32_t> indicePadParams{false, 0, 0, 0};
        AscendC::DataCopyExtParams indexCopyParams{1, 1 * 4, 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位

        AscendC::DataCopyPad(valuesGm[progress], tempValue, copyParams); // 从GM->VECIN搬运40Bytes
        AscendC::DataCopyPad(indiceGm[progress], tempIndex, copyParams); // 从GM->VECIN搬运40Bytes
        // for (int i = 0; i < 4; i++) {
        //     printf("max %d : %f\n", tempIndex.GetValue(i), tempValue.GetValue(i));
        // }
    }

private:
    int32_t coreDataNum;
    uint32_t totalLength;
    uint32_t globalIndex;
    uint32_t resLength;
    uint32_t stride;
    uint32_t iterStep;
    uint32_t nIter;
    int64_t axes;
    uint32_t dims;
    uint32_t axesDim;
    uint32_t loopCount;
    AscendC::TPipe *pipe;
    AscendC::GlobalTensor<DTYPE_X> xGm;
    AscendC::GlobalTensor<DTYPE_INDICE> indiceGm;
    AscendC::GlobalTensor<DTYPE_VALUES> valuesGm;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueue8X,inQueueFullX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueValues, outQueueIndice;
    TBuf<QuePosition::VECCALC>CastBuf1, CastBuf2, CastBuf3, computeBuf;
    TBuf<QuePosition::VECCALC>WorkQueue, xFullBuf;
    TBuf<QuePosition::VECCALC>Queue_tempIndex, Queue_tempValues, offsetLocal;
    LocalTensor<int32_t> tempIndex;
    LocalTensor<DTYPE_X> tempValue, xbuf;
    AscendC::LocalTensor<uint32_t> srcOffsetLocal;
};

template<typename T>class KernelArgMaxValueHalf{
public:
    __aicore__ inline KernelArgMaxValueHalf(){}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR indice, GM_ADDR values,
     uint32_t totalLength, uint32_t resLength, uint32_t stride,
     uint32_t iterStep, uint32_t nIter, int64_t axes, uint32_t dims, uint32_t axesDim,
     uint32_t smallDataNum,uint32_t bigDataNum,uint32_t tailBlockNum, TPipe* pipeIn
    ){
        this->totalLength = totalLength;
        this->resLength = resLength;
        this->stride = stride;
        this->iterStep = iterStep;
        this->nIter = nIter;
        this->axes = axes;
        this->dims = dims;
        this->axesDim = axesDim;
        pipe = pipeIn; 

        uint32_t coreNum = AscendC::GetBlockIdx();
        uint32_t globalIndex = bigDataNum * coreNum;
        if (coreNum < tailBlockNum) {
            this->coreDataNum = bigDataNum;
        } else {
            this->coreDataNum = smallDataNum;
            globalIndex -= (bigDataNum - smallDataNum)*(coreNum - tailBlockNum);
        }
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x, this->totalLength *sizeof(DTYPE_X));
        indiceGm.SetGlobalBuffer((__gm__ DTYPE_INDICE*)indice + globalIndex, this->coreDataNum * sizeof(DTYPE_INDICE));
        valuesGm.SetGlobalBuffer((__gm__ DTYPE_VALUES*)values + globalIndex, this->coreDataNum * sizeof(DTYPE_VALUES));
        this->globalIndex = globalIndex;

        this->loopCount = this->coreDataNum;
        pipe->InitBuffer(Queue_tempIndex,  2 * sizeof(int32_t));
        pipe->InitBuffer(Queue_tempValues, 2 * sizeof(DTYPE_X));

        pipe->InitBuffer(inQueueX, BUFFER_NUM, this->axesDim*16*sizeof(DTYPE_X));

        // pipe->InitBuffer(outQueueValues, BUFFER_NUM, this->stride*sizeof(DTYPE_X));
        // pipe->InitBuffer(outQueueIndice, BUFFER_NUM, this->stride*sizeof(int32_t));
        pipe->InitBuffer(computeBuf, 2*sizeof(DTYPE_X));
        pipe->InitBuffer(WorkQueue, this->axesDim*16*sizeof(DTYPE_X));
        // pipe->InitBuffer(offsetLocal, this->axesDim*16*sizeof(uint32_t));
        // pipe->InitBuffer(xFullBuf, this->axesDim*16*sizeof(DTYPE_X));

        this->tempIndex = Queue_tempIndex.Get<int32_t>();
        // this->xbuf = xFullBuf.Get<DTYPE_X>();
        // Duplicate(this->tempIndex, int32_t(0), this->stride);

        this->tempValue = Queue_tempValues.Get<DTYPE_X>();
    }
    // __aicore__ inline void process() {
    //     for (int i = 0; i < loopCount; i++) {
    //         for (int j = 0; j < stride; j++) {
    //             copyIn(i, j);
    //             compute(i, j);
    //         }
    //         copyOut(i);
    //     }
    // }
    __aicore__ inline void process() {
        for (int i = globalIndex; i < globalIndex + loopCount; i++) {
            printf("i: %d\n",i);
            copyIn(i/stride, i%stride);
            compute(i);
            copyOut(i-globalIndex);
        }
    }
    // __aicore__ void copyIn(int32_t progress, int now) {
    //     // LocalTensor<DTYPE_X>xLocal = inQueueX.AllocTensor<DTYPE_X>();
    //     LocalTensor<DTYPE_X>x16Local = inQueue16X.AllocTensor<DTYPE_X>();
    //     // TODO: check 是否需要-1
    //     //printf("stride: %d\n", this->stride*4);
    //     AscendC::DataCopyExtParams copyParams{static_cast<uint16_t>(this->axesDim),1*sizeof(half), (this->stride-1) * 2, 0,0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
    //     AscendC::DataCopyPadExtParams<DTYPE_X> padParams{false, 0, 0, 0};
    //     AscendC::DataCopyPad(x16Local, xGm[progress*iterStep + now], copyParams, padParams); // 从GM->VECIN搬运40Bytes
    //     // for (int i = 0; i < this->axesDim; i++) {
    //     //     //printf("xlocal i:%d value:%f\n",i,x16Local.GetValue(i*16));
    //     //     xLocal.SetValue(i, x16Local.GetValue(i*16));
    //     // }
    //     // printf("xLocal: %f %f %f \n",xLocal.GetValue(0), xLocal.GetValue(1), Local.GetValue(1));
    //     inQueue16X.EnQue<DTYPE_X>(x16Local);
    //     // inQueue16X.FreeTensor(x16Local);

    // }
    // 第一次就拷贝完整的元素

    __aicore__ inline void copyIn(int progress, int now) {
        LocalTensor<DTYPE_X>xLocal = inQueueX.AllocTensor<DTYPE_X>();
        AscendC::DataCopyExtParams copyParams{static_cast<uint16_t>(this->axesDim),  2, (this->stride-1) * 2, 0,0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPadExtParams<DTYPE_X> padParams{false, 0, 0, 0};
        // int x = process - process/stride;
        AscendC::DataCopyPad(xLocal, xGm[progress*iterStep + now], copyParams, padParams); // 从GM->VECIN搬运40Bytes
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void compute(int progress) {
        AscendC::LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        AscendC::LocalTensor<DTYPE_X> dstLocal = computeBuf.Get<DTYPE_X>();
        AscendC::LocalTensor<DTYPE_X> workLocal = WorkQueue.Get<DTYPE_X>();
        AscendC::ReduceMax(dstLocal ,xLocal, workLocal, this->axesDim*16, true);
        DTYPE_X maxValue = dstLocal.GetValue(0);
        DTYPE_X maxIndex = dstLocal.GetValue(1);
        uint32_t realIndex = *reinterpret_cast<uint32_t *>(&maxIndex);
        realIndex /= 16;
        tempIndex.SetValue(0, realIndex);
        tempValue.SetValue(0, maxValue);
        inQueueX.FreeTensor(xLocal);
    }

    __aicore__ void copyOut(int32_t progress) {
        AscendC::DataCopyExtParams copyParams{1, 1 * 2, 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        // AscendC::DataCopyPadExtParams<DTYPE_X> padParams{false, 0, 0, 0};
        // AscendC::DataCopyPadExtParams<int32_t> indicePadParams{false, 0, 0, 0};
        AscendC::DataCopyExtParams indexCopyParams{1, 1 * 4, 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位

        AscendC::DataCopyPad(valuesGm[progress], tempValue, copyParams); // 从GM->VECIN搬运40Bytes
        AscendC::DataCopyPad(indiceGm[progress], tempIndex, copyParams); // 从GM->VECIN搬运40Bytes
        // for (int i = 0; i < 4; i++) {
        //     printf("max %d : %f\n", tempIndex.GetValue(i), tempValue.GetValue(i));
        // }
    }

private:
    int32_t coreDataNum;
    uint32_t totalLength,globalIndex;
    uint32_t resLength;
    uint32_t stride;
    uint32_t iterStep;
    uint32_t nIter;
    int64_t axes;
    uint32_t dims;
    uint32_t axesDim;
    uint32_t loopCount;
    AscendC::TPipe* pipe;
    AscendC::GlobalTensor<DTYPE_X> xGm;
    AscendC::GlobalTensor<DTYPE_INDICE> indiceGm;
    AscendC::GlobalTensor<DTYPE_VALUES> valuesGm;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueue16X,inQueueFullX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueValues, outQueueIndice;
    TBuf<QuePosition::VECCALC>CastBuf1, CastBuf2, CastBuf3, computeBuf;
    TBuf<QuePosition::VECCALC>WorkQueue,xFullBuf,offsetLocal;
    TBuf<QuePosition::VECCALC>Queue_tempIndex, Queue_tempValues;
    LocalTensor<int32_t> tempIndex;
    LocalTensor<DTYPE_X> tempValue, xbuf;
    AscendC::LocalTensor<uint32_t> srcOffsetLocal;

};

template<typename T>
class KernelArgMaxValueInt32{
public:
    __aicore__ inline KernelArgMaxValueInt32(){}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR indice, GM_ADDR values,
     uint32_t totalLength, uint32_t resLength, uint32_t stride,
     uint32_t iterStep, uint32_t nIter, int64_t axes, uint32_t dims, uint32_t axesDim,
     uint32_t smallDataNum,uint32_t bigDataNum,uint32_t tailBlockNum, TPipe* pipeIn
    ){
        this->totalLength = totalLength;
        this->resLength = resLength;
        this->stride = stride;
        this->iterStep = iterStep;
        this->nIter = nIter;
        this->axes = axes;
        this->dims = dims;
        this->axesDim = axesDim;
        pipe = pipeIn; 

        uint32_t coreNum = AscendC::GetBlockIdx();
        uint32_t globalIndex = bigDataNum * coreNum;
        if (coreNum < tailBlockNum) {
            this->coreDataNum = bigDataNum;
        } else {
            this->coreDataNum = smallDataNum;
            globalIndex -= (bigDataNum - smallDataNum)*(coreNum - tailBlockNum);
        }
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + globalIndex * iterStep, this->coreDataNum * iterStep *sizeof(DTYPE_X));
        indiceGm.SetGlobalBuffer((__gm__ DTYPE_INDICE*)indice + globalIndex * stride, this->coreDataNum * iterStep *sizeof(DTYPE_INDICE));
        valuesGm.SetGlobalBuffer((__gm__ DTYPE_VALUES*)values + globalIndex * stride, this->coreDataNum * iterStep *sizeof(DTYPE_VALUES));
        
        this->loopCount = this->coreDataNum;
        pipe->InitBuffer(Queue_tempIndex,  this->stride * sizeof(int32_t));
        pipe->InitBuffer(Queue_tempValues, this->stride * sizeof(DTYPE_X));

        pipe->InitBuffer(inQueueX, BUFFER_NUM, this->axesDim*sizeof(DTYPE_X));
        pipe->InitBuffer(inQueue8X, BUFFER_NUM, this->axesDim*8*sizeof(DTYPE_X));

        pipe->InitBuffer(outQueueValues, BUFFER_NUM, this->stride*sizeof(DTYPE_X));
        pipe->InitBuffer(outQueueIndice, BUFFER_NUM, this->stride*sizeof(int32_t));
        pipe->InitBuffer(computeBuf, this->axesDim*sizeof(DTYPE_X));
        pipe->InitBuffer(WorkQueue, 8*this->axesDim*sizeof(float));
        pipe->InitBuffer(CastBuf1, 8*this->axesDim*sizeof(float));
        this->tempIndex = Queue_tempIndex.Get<int32_t>();
        Duplicate(this->tempIndex, int(0), this->stride);

        this->tempValue = Queue_tempValues.Get<DTYPE_X>();
        Duplicate(this->tempValue, DTYPE_X(0),this->stride);
    }
    __aicore__ inline void process() {
        for (int i = 0; i < loopCount; i++) {
            for (int j = 0; j < stride; j++) {
                copyIn(i, j);
                compute(i, j);
            }
            copyOut(i);
        }
    }
    __aicore__ void copyIn(int32_t progress, int now) {
        // LocalTensor<DTYPE_X>xLocal = inQueueX.AllocTensor<DTYPE_X>();
        LocalTensor<DTYPE_X>x8Local = inQueue8X.AllocTensor<DTYPE_X>();
        // TODO: check 是否需要-1
        //printf("stride: %d\n", this->stride*4);
        AscendC::DataCopyExtParams copyParams{static_cast<uint16_t>(this->axesDim),1*sizeof(float), (this->stride-1) * 4, 0,0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPadExtParams<DTYPE_X> padParams{false, 0, 0, 0};
        AscendC::DataCopyPad(x8Local, xGm[progress*iterStep + now], copyParams, padParams); // 从GM->VECIN搬运40Bytes
        // for (int i = 0; i < this->axesDim; i++) {
        //     //printf("xlocal i:%d value:%f\n",i,x8Local.GetValue(i*8));
        //     xLocal.SetValue(i, x8Local.GetValue(i*8));
        // }
        // printf("xLocal: %f %f %f \n",xLocal.GetValue(0), xLocal.GetValue(1), Local.GetValue(1));
        // inQueueX.EnQue<DTYPE_X>(xLocal);
        inQueue8X.EnQue<DTYPE_X>(x8Local);

    }

    __aicore__ void compute(int32_t progress, int now) {
        AscendC::LocalTensor<DTYPE_X> xLocal = inQueue8X.DeQue<DTYPE_X>();
        AscendC::LocalTensor<float> dstLocal = computeBuf.Get<float>();
        AscendC::LocalTensor<float> workLocal = WorkQueue.Get<float>();
        auto castBuf = CastBuf1.Get<float>();
        Cast(castBuf, xLocal, RoundMode::CAST_NONE,8*this->axesDim);
        AscendC::ReduceMax(dstLocal ,castBuf, workLocal, 8*this->axesDim ,true);
        float maxValue = dstLocal.GetValue(0);
        float maxIndex = dstLocal.GetValue(1);
        uint32_t realIndex = *reinterpret_cast<uint32_t *>(&maxIndex);
        int32_t realMax = static_cast<int32_t>(maxValue);
        //printf("maxValue:%f maxIndex:%d realMax:%d\n", maxValue, realIndex, realMax);
        realIndex /= 8;
        tempIndex.SetValue(now, realIndex);
        tempValue.SetValue(now, realMax);
        inQueue8X.FreeTensor(xLocal);

    }
    __aicore__ void copyOut(int32_t progress) {
        AscendC::DataCopyExtParams copyParams{1, this->stride * 4, 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        // AscendC::DataCopyPadExtParams<DTYPE_X> padParams{false, 0, 0, 0};
        // AscendC::DataCopyPadExtParams<int32_t> indicePadParams{false, 0, 0, 0};
        AscendC::DataCopyExtParams indexCopyParams{1, this->stride * 4, 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位

        AscendC::DataCopyPad(valuesGm[progress*stride], tempValue, copyParams); // 从GM->VECIN搬运40Bytes
        AscendC::DataCopyPad(indiceGm[progress*stride], tempIndex, copyParams); // 从GM->VECIN搬运40Bytes
        // for (int i = 0; i < 4; i++) {
        //     printf("%d : %f\n", tempIndex.GetValue(i), tempValue.GetValue(i));
        // }
    }

private:
    int32_t coreDataNum;
    uint32_t totalLength;
    uint32_t resLength;
    uint32_t stride;
    uint32_t iterStep;
    uint32_t nIter;
    int64_t axes;
    uint32_t dims;
    uint32_t axesDim;
    uint32_t loopCount;
    AscendC::TPipe* pipe;
    AscendC::GlobalTensor<DTYPE_X> xGm;
    AscendC::GlobalTensor<DTYPE_INDICE> indiceGm;
    AscendC::GlobalTensor<DTYPE_VALUES> valuesGm;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueue8X;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueValues, outQueueIndice;
    TBuf<QuePosition::VECCALC>CastBuf1, CastBuf2, CastBuf3, computeBuf;
    TBuf<QuePosition::VECCALC>WorkQueue;
    TBuf<QuePosition::VECCALC>Queue_tempIndex, Queue_tempValues;
    LocalTensor<int32_t> tempIndex;
    LocalTensor<DTYPE_X> tempValue;
};

class KernelArgMaxValueInt8{
public:
    __aicore__ inline KernelArgMaxValueInt8(){}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR indice, GM_ADDR values,
     uint32_t totalLength, uint32_t resLength, uint32_t stride,
     uint32_t iterStep, uint32_t nIter, int64_t axes, uint32_t dims, uint32_t axesDim,
     uint32_t smallDataNum,uint32_t bigDataNum,uint32_t tailBlockNum, TPipe* pipeIn
    ){
        this->totalLength = totalLength;
        this->resLength = resLength;
        this->stride = stride;
        this->iterStep = iterStep;
        this->nIter = nIter;
        this->axes = axes;
        this->dims = dims;
        this->axesDim = axesDim;
        pipe = pipeIn; 

        uint32_t coreNum = AscendC::GetBlockIdx();
        uint32_t globalIndex = bigDataNum * coreNum;
        if (coreNum < tailBlockNum) {
            this->coreDataNum = bigDataNum;
        } else {
            this->coreDataNum = smallDataNum;
            globalIndex -= (bigDataNum - smallDataNum)*(coreNum - tailBlockNum);
        }
        // printf("gc: %d %d\n", globalIndex, coreDataNum);
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + globalIndex, this->totalLength *sizeof(DTYPE_X));
        indiceGm.SetGlobalBuffer((__gm__ DTYPE_INDICE*)indice + globalIndex, this->coreDataNum *sizeof(DTYPE_INDICE));
        valuesGm.SetGlobalBuffer((__gm__ DTYPE_VALUES*)values + globalIndex, this->coreDataNum *sizeof(DTYPE_VALUES));
        
        this->loopCount = this->coreDataNum;
        pipe->InitBuffer(Queue_tempIndex,  sizeof(int32_t));
        pipe->InitBuffer(Queue_tempValues, sizeof(DTYPE_X));

        pipe->InitBuffer(inQueueX, BUFFER_NUM, this->axesDim*sizeof(DTYPE_X));
        pipe->InitBuffer(inQueue32X, BUFFER_NUM, this->axesDim*32*sizeof(DTYPE_X));
        

        // pipe.InitBuffer(outQueueValues, BUFFER_NUM, this->stride*sizeof(DTYPE_X));
        // pipe.InitBuffer(outQueueIndice, BUFFER_NUM, this->stride*sizeof(int32_t));
        pipe->InitBuffer(computeBuf, this->axesDim*sizeof(DTYPE_X));
        pipe->InitBuffer(WorkQueue, 32*this->axesDim*sizeof(half));

        pipe->InitBuffer(CastBuf1, 32*this->axesDim*sizeof(half));
        pipe->InitBuffer(CastBuf2, 32*this->axesDim*sizeof(half));


        this->tempIndex = Queue_tempIndex.Get<int32_t>();
        Duplicate(this->tempIndex, int32_t(0), 1);

        this->tempValue = Queue_tempValues.Get<DTYPE_X>();
    }
    __aicore__ inline void process() {
        for (int i = 0; i < coreDataNum; i++) {
                copyIn(i);
                compute(i);
                copyOut(i);
        }
    }
    __aicore__ void copyIn(int32_t progress) {
        // LocalTensor<DTYPE_X>xLocal = inQueueX.AllocTensor<DTYPE_X>();
        LocalTensor<DTYPE_X>x32Local = inQueue32X.AllocTensor<DTYPE_X>();
        // TODO: check 是否需要-1
        // printf("stride: %d\n", this->stride*1);
        AscendC::DataCopyExtParams copyParams{static_cast<uint16_t>(this->axesDim),1*1/* uint8 */, (this->stride-1) * 1, 0,0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        AscendC::DataCopyPadExtParams<DTYPE_X> padParams{false, 0, 0, 0};
        AscendC::DataCopyPad(x32Local, xGm[progress], copyParams, padParams); // 从GM->VECIN搬运40Bytes
        // for (int i = 0; i < this->axesDim; i++) {
        //     // printf("xlocal i:%d value:%f\n",i,x32Local.GetValue(i*32));
        //     xLocal.SetValue(i, x32Local.GetValue(i*32));
        // }
        // printf("xLocal: %f %f %f \n",xLocal.GetValue(0), xLocal.GetValue(1), Local.GetValue(1));
        // inQueueX.EnQue<DTYPE_X>(xLocal);
        inQueue32X.EnQue<DTYPE_X>(x32Local);

    }

    __aicore__ void compute(int32_t progress) {
        AscendC::LocalTensor<DTYPE_X> xLocal = inQueue32X.DeQue<DTYPE_X>();
        AscendC::LocalTensor<DTYPE_X> dstLocal = computeBuf.Get<DTYPE_X>();
        auto workLocal = WorkQueue.Get<float16_t>();
        auto castbuf = CastBuf1.Get<float16_t>();
        auto castbufout = CastBuf2.Get<float16_t>();
        auto intx = xLocal.ReinterpretCast<uint8_t>();
        auto intdst = dstLocal.ReinterpretCast<uint8_t>();
        

        AscendC::Cast(castbuf, intx, RoundMode::CAST_NONE,32* this->axesDim); 

        AscendC::ReduceMax(castbufout ,castbuf, workLocal,32* this->axesDim ,true);
        // for (int i = 0; i <= 1; i++) {
        //     printf("%f ",castbuf.GetValue(i));
        // }
        // printf("\n");
        Cast(intdst, castbufout, RoundMode::CAST_TRUNC, 1); 
        // DataSyncBarrier<MemDsbT::ALL>();
        // for (int i = 0; i < 1; i++)
        // {
        //     printf("[CAST] xLocal,castbuf,castbufout,dstLocal: %d %f %f %d \n",xLocal.GetValue(i), castbuf.GetValue(i), castbufout.GetValue(i), dstLocal.GetValue(i));
        // }
        

        DTYPE_X maxValue = dstLocal.GetValue(0);
        half maxIndex = castbufout.GetValue(1);
        uint32_t realIndex = *reinterpret_cast<uint16_t *>(&maxIndex);
        // printf("maxValue:%f maxIndex:%d\n", maxValue, realIndex);
        realIndex /= 32;
        tempIndex.SetValue(0, realIndex);
        tempValue.SetValue(0, maxValue);
        inQueue32X.FreeTensor(xLocal);

    }
    __aicore__ void copyOut(int32_t progress) {
        AscendC::DataCopyExtParams copyParams{1, 1, 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
        // AscendC::DataCopyPadExtParams<DTYPE_X> padParams{false, 0, 0, 0};
        // AscendC::DataCopyPadExtParams<int32_t> indicePadParams{false, 0, 0, 0};
        AscendC::DataCopyExtParams indexCopyParams{1, 4, 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位

        AscendC::DataCopyPad(valuesGm[progress], tempValue, copyParams); // 从GM->VECIN搬运40Bytes
        AscendC::DataCopyPad(indiceGm[progress], tempIndex,  indexCopyParams); // 从GM->VECIN搬运40Bytes
        // for (int i = 0; i < 4; i++) {
            // printf("%d : %d\n", tempIndex.GetValue(0), tempValue.GetValue(0));
        // }
    }

private:
    int32_t coreDataNum;
    uint32_t totalLength;
    uint32_t resLength;
    uint32_t stride;
    uint32_t iterStep;
    uint32_t nIter;
    int64_t axes;
    uint32_t dims;
    uint32_t axesDim;
    uint32_t loopCount;
    AscendC::TPipe* pipe;
    AscendC::GlobalTensor<DTYPE_X> xGm;
    AscendC::GlobalTensor<DTYPE_INDICE> indiceGm;
    AscendC::GlobalTensor<DTYPE_VALUES> valuesGm;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueue32X;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueValues, outQueueIndice;
    TBuf<QuePosition::VECCALC>CastBuf1, CastBuf2, CastBuf3, computeBuf;
    TBuf<QuePosition::VECCALC>WorkQueue;
    TBuf<QuePosition::VECCALC>Queue_tempIndex, Queue_tempValues;
    LocalTensor<int32_t> tempIndex;
    LocalTensor<DTYPE_X> tempValue;

};

extern "C" __global__ __aicore__ void arg_max_with_value(GM_ADDR x, GM_ADDR indice, GM_ADDR values, GM_ADDR workspace, GM_ADDR tiling) {

    TPipe pipe;
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); // 增加这一行
    if (TILING_KEY_IS(1)) {
        if constexpr(std::is_same_v<DTYPE_X, float>){
                KernelArgMaxValuePerf<DTYPE_X>op;
                auto SmartTiling= reinterpret_cast<__gm__ uint32_t*>(tiling);
                uint32_t magicNum =  *SmartTiling;
                uint32_t res = magicNum/1000;
                uint32_t smallDataNum =  *(SmartTiling + 1);
                uint32_t tail =  *(SmartTiling + 2);
                // op.Init(x, indice, values,
                // tiling_data.axesDim, tiling_data.smallDataNum, tiling_data.bigDataNum, tiling_data.tailBlockNum, &pipe);
                op.Init(x, indice, values, axesDim, smallDataNum, smallDataNum+1, tail, &pipe);

                op.process();
                return;
        }
    }
      else if (TILING_KEY_IS(0)) {
        GET_TILING_DATA(tiling_data, tiling);
        if constexpr(std::is_same_v<DTYPE_X, float>){
                if (tiling_data.axesDim <= 64) {
                GET_TILING_DATA(tiling_data, tiling);    
                KernelArgMaxValueSpec<DTYPE_X>op;
                    op.Init(x, indice, values,
                    tiling_data.axesDim, tiling_data.smallDataNum, tiling_data.bigDataNum, tiling_data.tailBlockNum, &pipe);
                    op.process1();
                    return;
                }
            }
        if constexpr(std::is_same_v<DTYPE_X, half>){
                KernelArgMaxValueHalf<DTYPE_X>op;
                op.Init(x, indice, values, tiling_data.totalLength, tiling_data.resLength, tiling_data.stride, tiling_data.iterStep, tiling_data.nIter, tiling_data.axes, tiling_data.dims,
                tiling_data.axesDim, tiling_data.smallDataNum, tiling_data.bigDataNum, tiling_data.tailBlockNum, &pipe);
                op.process();
                return;        
            }  else if constexpr(std::is_same_v<DTYPE_X, int32_t>){
                KernelArgMaxValueInt32<int32_t> op;
                op.Init(x, indice, values, tiling_data.totalLength, tiling_data.resLength, tiling_data.stride, tiling_data.iterStep, tiling_data.nIter, tiling_data.axes, tiling_data.dims,
                tiling_data.axesDim, tiling_data.smallDataNum, tiling_data.bigDataNum, tiling_data.tailBlockNum, &pipe);
                op.process();
                return;        
            }else{
                KernelArgMaxValueInt8 op;
                printf("-----------------------Processing uint8_t\n");
                op.Init(x, indice, values, tiling_data.totalLength, tiling_data.resLength, tiling_data.stride, tiling_data.iterStep, tiling_data.nIter, tiling_data.axes, tiling_data.dims,
                tiling_data.axesDim, tiling_data.smallDataNum, tiling_data.bigDataNum, tiling_data.tailBlockNum, &pipe);
                op.process();
                return;  
            }
      }
}
